from flask import Flask, request, jsonify
from flask_cors import CORS
import os
from dashscope import Generation
import json

app = Flask(__name__)
CORS(app)

# 定义工具列表
tools = [
    {
        "type": "function",
        "function": {
            "name": "add",
            "description": "用于执行两个数字的加法运算",
            "parameters": {
                "type": "object",
                "properties": {
                    "num1": {
                        "type": "number",
                        "description": "第一个数字"
                    },
                    "num2": {
                        "type": "number",
                        "description": "第二个数字"
                    }
                },
                "required": ["num1", "num2"]
            }
        }
    }
]

# 加法函数实现
def add(num1, num2):
    result = num1 + num2
    print(f"Result of {num1} + {num2} is {result}")
    return f"{num1} + {num2} = {result}"


@app.route('/api/chat', methods=['POST'])
def chat():
    data = request.json
    message = data.get('message')
    
    messages = [
        {'role': 'system', 'content': 'You are a helpful assistant.'},
        {'role': 'user', 'content': message}
    ]
    
    try:
        response = Generation.call(
            api_key=os.environ.get('DASH_SCOPE_API_KEY'),
            model="qwen-plus",
            messages=messages,
            tools=tools,
            result_format="message"
        )
        
        if response.status_code == 200:
            assistant_output = response.output.choices[0].message
            
            # 检查是否需要调用工具
            if 'tool_calls' not in assistant_output:
                return jsonify({
                    'success': True,
                    'response': assistant_output.content
                })
            
            # 处理工具调用
            if assistant_output.tool_calls[0]['function']['name'] == 'add':
                args = json.loads(assistant_output.tool_calls[0]['function']['arguments'])
                tool_response = add(args['num1'], args['num2'])
                
                # 将工具响应添加到消息列表
                messages.append(assistant_output)
                messages.append({
                    'name': 'add',
                    'role': 'tool',
                    'content': tool_response
                })
                
                # 再次调用模型进行总结
                final_response = Generation.call(
                    api_key="sk-d418cbb18a4b44eb821507683be1d6fe",
                    model="qwen-plus",
                    messages=messages,
                    tools=tools,
                    result_format="message"
                )
                
                return jsonify({
                    'success': True,
                    'response': final_response.output.choices[0].message.content
                })
            
            return jsonify({
                'success': True,
                'response': assistant_output.content
            })
        else:
            return jsonify({
                'success': False,
                'error': f"Error: {response.code} - {response.message}"
            }), response.status_code
            
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

if __name__ == '__main__':
    app.run(debug=True)